from PIL import Image
import numpy as np

class ROP_single_clinic():
    def __init__(self,image_resolution=(224,224),threshold=0.5,bbox_expand=20):
        super(ROP_single_clinic, self).__init__()
        
        self.image_resolution=image_resolution
        
        self.threshold=threshold
        self.bbox_expand=bbox_expand

    def get_lesion_location(self,segmentation_mask_path):
        # open the image
        segmentation_mask = Image.open(segmentation_mask_path).resize(self.image_resolution)
        segmentation_mask = np.array(segmentation_mask)
        # Get the lesion location
        segmentation_mask = segmentation_mask
        lesion_location = np.where(segmentation_mask>self.threshold)
        
        # Get the bounding box
        xmin = np.min(lesion_location[1])
        xmax = np.max(lesion_location[1])
        ymin = np.min(lesion_location[0])
        ymax = np.max(lesion_location[0])
        
        # expand the bounding box
        xmin = max(0,xmin-self.bbox_expand)
        ymin = max(0,ymin-self.bbox_expand)
        xmax = min(segmentation_mask.shape[1],xmax+self.bbox_expand)
        ymax = min(segmentation_mask.shape[0],ymax+self.bbox_expand)
        
        return [[xmin,xmax,ymin,ymax]]